from multiprocessing import shared_memory
import time
import win32gui
import win32con

import torch.nn.functional as F
import torch
import numpy as np
import cv2
from PIL import Image
import torchvision

from torch_v2 import LeNet as Net
from grabscreen import grab_screen

if __name__ == '__main__':

    network = Net()
    network.eval()

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    network.load_state_dict(torch.load('./model/net_008.pth', map_location=device))
    network = network.to(device)

    # transforms = get_train_transforms(input_size=(28, 28))
    with torch.no_grad():

        mode = 'screenshot'
        if mode == 'screenshot':
            # x1, y1, x2, y2 = int(1920 / 2 - 416 / 2), int(1080 / 2 - 416 / 2), int(1920 / 2 + 416 / 2), int(
            #     1080 / 2 + 416 / 2)

            # x1, y1 = int(958 - 48 / 2), int(1003 - 48 / 2)
            # x2, y2 = int(958 + 48 / 2), int(1003 + 48 / 2)

            x1, y1 = int(958 - 32 / 2), int(1006 - 32 / 2)
            x2, y2 = int(958 + 32 / 2), int(1006 + 32 / 2)

            # x1, y1 = int(965 - 16 / 2), int(1006 - 32 / 2)
            # x2, y2 = int(965 + 16 / 2), int(1006 + 32 / 2)

            fps = 0.0

            view_cv = False  # 是否预览
            view_cv = True  # 是否预览
            view_size = 200

            if view_cv:
                cv2.namedWindow('csgo_detect', cv2.WINDOW_NORMAL)
                cv2.resizeWindow('csgo_detect', view_size, view_size)

            # listener = pynput.mouse.Listener(
            #     # on_move=on_move,
            #     on_click=on_click,
            #     # on_scroll=on_scroll
            # )
            # listener.start()
            # try:  # 判断姿势
            #     share_posture = shared_memory.ShareableList([1], name='share_posture')
            # except FileExistsError:
            #     share_posture = shared_memory.ShareableList(name='share_posture')

            # try:  # 判断是否处于瞄准状态
            #     share_lock_mode = shared_memory.ShareableList([0], name='share_lock_mode')
            # except FileExistsError:
            #     share_lock_mode = shared_memory.ShareableList(name='share_lock_mode')

            try:  # 判断是否处于开火状态
                share_fire_mode = shared_memory.ShareableList([0], name='share_fire_mode')
            except FileExistsError:
                share_fire_mode = shared_memory.ShareableList(name='share_fire_mode')

            # pre_transforms = torchvision.transforms.Compose([
            #     torchvision.transforms.ToPILImage(),
            #     # torchvision.transforms.SSDCropping(),
            #     # torchvision.transforms.Resize(size=(256, 256)),
            #     torchvision.transforms.Resize(size=(28, 28)),
            #     # torchvision.transforms.ColorJitter(),
            #     torchvision.transforms.Grayscale(num_output_channels=1),  # 转灰度图
            #     torchvision.transforms.ToTensor(),
            #     # torchvision.transforms.RandomHorizontalFlip(),
            #     # torchvision.transforms.Normalization(),
            #     # torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),  # 归一化
            #     # torchvision.transforms.AssignGTtoDefaultBox()
            # ])

            pos_cls = ['lie', 'squat', 'stand']

            lock_mode = False

            last_img = None
            last_ts = time.time()

            fire_delay = 0.15  # 最小持续开火时间，越大则处于开火压枪模式概率越大

            f_diff = np.frompyfunc(lambda a, b: a - b if a > b else b - a, 2, 1)

            while True:

                t1 = time.time()
                # time.sleep(0.08)  # 延时

                img0 = grab_screen(region=(x1, y1, x2, y2))

                if view_cv:
                    r_img0 = cv2.resize(img0, dsize=(view_size, view_size))
                    r_img0 = cv2.cvtColor(r_img0, cv2.COLOR_BGR2GRAY)
                    ret, r_img0 = cv2.threshold(r_img0, 200, 255, cv2.THRESH_BINARY)

                img0 = cv2.cvtColor(img0, cv2.COLOR_BGR2GRAY)
                ret, img0 = cv2.threshold(img0, 200, 255, cv2.THRESH_BINARY)

                current_ts = time.time()
                if img0.sum() > 20000:  # 判断是否有数字（是否可以开火）
                    if last_img is None:
                        last_img = img0

                    if (d := f_diff(img0, last_img).sum()) > 20000:  # 是否处于开火状态
                        share_fire_mode[0] = 1
                        # print(d)
                        last_ts = current_ts
                    else:
                        if current_ts - last_ts < fire_delay:
                            pass
                        else:
                            share_fire_mode[0] = 0

                    last_img = img0
                else:
                    if current_ts - last_ts < fire_delay:
                        pass
                    else:
                        share_fire_mode[0] = 0

                if view_cv:

                    fps = round((fps + (1. / (time.time() - t1 + 0.00001))) / 2)
                    r_img0 = cv2.putText(r_img0, f"FPS {fps}", (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,), 2)
                    r_img0 = cv2.putText(r_img0, f"{share_fire_mode[0]}", (0, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,),
                                         2)
                    # r_img0 = cv2.putText(r_img0, f"POS {pos_cls[posture.argmax()]}", (0, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
                    # r_img0 = cv2.putText(r_img0, f"Aim {share_lock_mode[0]}", (0, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

                    cv2.imshow('csgo_detect', r_img0)

                    hwnd = win32gui.FindWindow(None, 'csgo_detect')
                    CVRECT = cv2.getWindowImageRect('csgo_detect')
                    win32gui.SetWindowPos(hwnd, win32con.HWND_TOPMOST, 0, 0, 0, 0,
                                          win32con.SWP_NOMOVE | win32con.SWP_NOSIZE)

                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        cv2.destroyAllWindows()
                        break

# # 每跑完一次epoch测试一下准确率
# with torch.no_grad():
#     correct = 0
#     total = 0
#
#
#
#     images = images.to(device)
#     outputs = net(images)
#     # 取得分最高的那个类
#     _, predicted = torch.max(outputs.data, 1)
#
